/*
 *        (C) COPYRIGHT Ingenic Limited.
 *             ALL RIGHTS RESERVED
 *
 * File       : inference_yolov5s.cpp
 * Authors    : ljhong, cywang
 * Create Time: 2024-11-21:15:57:21
 * Description:
 * 
 */

#include "ingenic_mnni.h"
#include "common/common_utils.h"
#include "utils/imgproc.h"
#include "utils/postproc.h"

#include <cmath>
#include <errno.h>
#include <fstream>
#include <iostream>
#include <memory>
#include <stdio.h>
#include <stdlib.h>
#include <string>
#include <cstring>
#include <string.h>
#include <vector>
#include <algorithm>
#include <sys/stat.h>   
#include <sys/types.h>
#include <dirent.h>     
#include <unistd.h>     
#include <sstream>      
#include <iomanip>      
#include <numeric>  // 添加 numeric 头文件，用于 std::inner_product
#include "yolo_interface.h"
#include "xb_image_utils.h"
#include <chrono>  // 添加 chrono 头文件
#include <stdio.h>


// // #define IS_ALIGN_64(x) (((size_t)x) & 0x3F)
#ifdef URANUS_NV12
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::NV12;
#else
magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::RGBA;
#endif


using namespace std;
using namespace magik::uranus;
using namespace Detect;
using namespace XBUtils;

typedef unsigned char uint8_t;

// 辅助函数：获取文件夹中所有图片路径
vector<string> get_image_paths(const string& dir_path) {
    vector<string> image_paths;
    DIR* dir = opendir(dir_path.c_str());
    if (!dir) {
        perror("opendir failed");
        return image_paths;
    }
    struct dirent* entry;
    while ((entry = readdir(dir)) != nullptr) {
        string filename = entry->d_name;
        if (filename == "." || filename == "..") continue;
        // 简单判断图片格式（可根据实际需求扩展）
        if (filename.find(".jpg") != string::npos || filename.find(".png") != string::npos) {
            image_paths.push_back(dir_path + "/" + filename);
        }
    }
    closedir(dir);
    return image_paths;
}

// 辅助函数：创建文件夹（递归）
bool create_directory(const string& path) {
    struct stat st;
    if (stat(path.c_str(), &st) == 0) {
        return S_ISDIR(st.st_mode); // 已存在且是目录
    }
    // 递归创建父目录
    size_t pos = path.find_last_of('/');
    if (pos != string::npos) {
        string parent = path.substr(0, pos);
        if (!create_directory(parent)) return false;
    }
    return mkdir(path.c_str(), 0755) == 0;
}

int remove_directory(const char *path) {
    DIR *dir = opendir(path);
    if (!dir) return -1;
    
    struct dirent *entry;
    while ((entry = readdir(dir)) != NULL) {
        if (strcmp(entry->d_name, ".") == 0 || strcmp(entry->d_name, "..") == 0)
            continue;
            
        char full_path[PATH_MAX];
        snprintf(full_path, sizeof(full_path), "%s/%s", path, entry->d_name);
        
        if (entry->d_type == DT_DIR) {
            remove_directory(full_path);
        } else {
            unlink(full_path);
        }
    }
    closedir(dir);
    return rmdir(path);
}

int main(int argc, char **argv) {
    if (argc != 4) {
        printf("Usage: %s det_model_path input_image_dir output_dir\n", argv[0]);
        exit(0);
    }
    std::string det_model_path = argv[1];
    std::string input_dir = argv[2];
    std::string output_dir = argv[3];

    // 创建输出目录（如果存在则先删除）
    if (access(output_dir.c_str(), F_OK) == 0) {
        if (remove_directory(output_dir.c_str()) != 0) {
            perror("remove existing directory failed");
            return -1;
        }
    }

    if (!create_directory(output_dir)) {
        perror("create output directory failed");
        return -1;
    }

    string crop_dir = output_dir + "/crop";
    if (!create_directory(crop_dir)) {
        perror("create crop_dir directory failed");
        return -1;
    }

    string resize_dir = output_dir + "/resize224";
    if (!create_directory(resize_dir)) {
        perror("create resize_dir directory failed");
        return -1;
    }

    string yolo_input_dir = output_dir + "/yolo_input";
    if (!create_directory(yolo_input_dir)) {
        perror("create yolo_input_dir directory failed");
        return -1;
    }

    // 加载检测模型
    YOLOV5Config yolo5_config;
    yolo5_config.conf_thresh = 0.6f;       // 修改置信度阈值
    yolo5_config.nms_thres = 0.45f;        // 修改NMS阈值
    yolo5_config.classes = 1;              // 修改类别数
    yolo5_config.model_path = det_model_path; // 设置模型路径
    yolo5_config.input_h = 640;
    yolo5_config.input_w = 640;

    auto yolov5_interface = YOLOV5Interface(yolo5_config);
    yolov5_interface.loadModel();
    yolov5_interface.print_model_forward_memory();

    // 获取所有输入图片路径
    vector<string> image_paths = get_image_paths(input_dir);
    int total_images = image_paths.size();
    int detected_images = 0;
    vector<string> undetected_paths;

    // 新增：初始化时间统计变量
    chrono::duration<double> total_elapsed = chrono::duration<double>::zero();

    for (size_t i = 0; i < image_paths.size(); ++i) {
        const auto& image_path = image_paths[i];
        // 新增：记录单张图片处理开始时间
        auto image_start_time = chrono::system_clock::now();
        // 读取图片
        printf("\nProcessing image: %s \n", image_path.c_str());
        FrameInfo input_info;
        MyImg img;
        int ret = ImageUtils::read_img(input_info, img, image_path.c_str()); 
        if(ret != 0)  {
            printf("read image failed: %s\n", image_path.c_str());
            continue;
        }
        string filename = image_path.substr(image_path.find_last_of('/') + 1);
        string output_path = output_dir + "/" + filename;

        magik::uranus::Box_t tmp_box;
        tmp_box = {0, 0, static_cast<float>(input_info.width), static_cast<float>(input_info.height)};
        /******************prepare data***********************/
        magik::uranus::shape_t yolo_working_img_shape = {1, yolo5_config.input_h, yolo5_config.input_w, 4};
        magik::uranus::Tensor working_img(yolo_working_img_shape);

        ImageUtils::prepare_data(input_info,  working_img, tmp_box, yolo5_config.input_w, yolo5_config.input_h, DATA_FORMAT); 
        MyImg yolo_input_img;
        ImageUtils::tensor_to_img(working_img, yolo_input_img);
        string yolo_output_path = yolo_input_dir + "/" + filename;
        ImageUtils::saveImg(yolo_input_img, yolo_output_path.c_str(), 60, std::string());
        if(yolo_input_img.data != nullptr){
            ImageUtils::free_stb_mem(yolo_input_img.data);
        }
        std::vector<magik::uranus::ObjBbox_t> output_boxes;
        yolov5_interface.detect(working_img, input_info.width, input_info.height, output_boxes);

        // 打印检测结果
        int num_boxes = output_boxes.size();
        std::cout << "num_boxes: " << num_boxes << std::endl;
        if (num_boxes > 0) detected_images++;
        else undetected_paths.push_back(image_path);
        for (const auto& box : output_boxes) {
            // 打印检测框信息
            printf("box:   ");
            printf("x0: %d ", (int)box.box.x0);
            printf("y0: %d ", (int)box.box.y0);
            printf("x1: %d ", (int)box.box.x1);
            printf("y1: %d ", (int)box.box.y1);
            printf("score: %.2f ", box.score);
            printf("class: %d ", box.class_id);
            printf("\n");
            // 裁剪图片
            MyImg crop_img;
            ImageUtils::crop_image(img, box.box, crop_img);
            string crop_output_path = crop_dir + "/" + filename;
            ImageUtils::saveImg(crop_img, crop_output_path.c_str(), 60, std::string());
            // 调整图片大小，resize
            MyImg resize_img;
            ImageUtils::resize_image(crop_img, resize_img, 224, 224);
            string resize_output_path = resize_dir + "/" + filename;
            ImageUtils::saveImg(resize_img, resize_output_path.c_str(), 60, std::string());
            // 释放内存
            if(crop_img.data != nullptr){
                ImageUtils::free_stb_mem(crop_img.data);
            }
            if(resize_img.data != nullptr){
                ImageUtils::free_stb_mem(resize_img.data);
            }
        } 

        // 画/保存检测框
        ImageUtils::draw_boxes(img, output_boxes); 
        ImageUtils::saveImg(img, output_path.c_str(), 60, std::string());
        if(img.data != nullptr){
            ImageUtils::free_stb_mem(img.data);
        }
        if(input_info.data != nullptr){
            ImageUtils::free_stb_mem(input_info.data);
        }

        // 新增：计算单张图片处理耗时
        auto image_end_time = chrono::system_clock::now();
        chrono::duration<double> image_duration = image_end_time - image_start_time;
        total_elapsed += image_duration;

        // 新增：计算进度和剩余时间
        int current = i + 1;
        double progress = (total_images == 0) ? 0 : (static_cast<double>(current) / total_images) * 100;
        double avg_time_per_image = total_elapsed.count() / current;
        double remaining_time = avg_time_per_image * (total_images - current);

        // 输出进度信息（保留2位小数）
        printf("进度：%d/%d (%.2f%%) | 已用时间：%.2fs | 剩余时间：%.2fs\n",
               current, total_images,
               progress,
               total_elapsed.count(),
               remaining_time);
    }

    // 生成统计文档
    string stats_path = output_dir + "/statistics.txt";
    ofstream stats_file(stats_path);
    if (stats_file.is_open()) {
        stats_file << "总图片数: " << total_images << "\n";
        stats_file << "检测到猫脸的图片数: " << detected_images << "\n";
        stats_file << "检测到猫脸的比例: " << (total_images == 0 ? 0.0 : (double)detected_images / total_images * 100) << "%\n";
        stats_file << "未检测到猫脸的图片路径:\n";
        for (const auto& path : undetected_paths) {
            stats_file << path << "\n";
        }
        stats_file.close();
        printf("统计文档已保存至: %s\n", stats_path.c_str());
    } else {
        printf("无法创建统计文档\n");
    }

    yolov5_interface.releaseModel();
    return 0;
}